import csv
import random
import os

# 设置目标记录数量
TARGET_RECORDS = 200  # 可以自由调整这个数值

# 定义中国移动产品
cmcc_products = [
    "移动看家", "移动高清", "移动康养", "移动车家", "移动云盘",
    "139邮箱", "云手机", "云电脑", "行车卫士", "集团语音",
    "e企组网", "专线卫士", "和对讲", "云视讯", "千里眼",
    "移动ICT", "九天大模型", "移动云"
]

# 定义企业及其行业分类
companies_by_industry = {
    "通信": ["华为技术有限公司", "中兴通讯股份有限公司", "烽火通信科技股份有限公司"],
    "科技": ["中国电子集团", "阿里巴巴集团", "腾讯控股集团", "中国移动通信集团",
             "中国电信集团", "中国联合网络通信集团", "网易网络有限公司", "百度公司",
             "京东集团", "小米集团"],
    "金融": ["中国工商银行集团", "中国平安集团", "中国建设银行集团", "中国银行集团",
             "中国人寿保险集团", "中国农业银行集团"],
    "能源": ["国家电网公司", "中国石油集团", "中国石化集团", "中国海洋石油集团",
             "国家能源投资集团"],
    "基建": ["中国建筑集团", "中国交通建设集团", "中国铁路工程集团", "中国铁道建筑集团"],
    "教育": ["新东方教育集团", "好未来教育集团", "北京大学", "清华大学", "复旦大学"],
    "物流": ["顺丰控股集团", "阿里巴巴集团", "京东物流集团", "中通快递", "圆通速递"],
    "医疗": ["华润医药集团", "国药集团", "上海医药集团", "九州通医药集团"],
    "制造": ["比亚迪集团", "中国电子集团", "海尔集团", "格力电器", "美的集团"],
    "餐饮": ["海底捞餐饮集团", "百胜中国控股", "美团点评", "饿了么", "星巴克中国"],
    "政务": ["国家税务总局", "公安部", "国家卫生健康委员会", "教育部", "交通运输部"],
    "零售": ["沃尔玛中国", "永辉超市", "苏宁易购", "国美零售", "拼多多"]
}

# 定义竞品映射
competitor_mapping = {
    "阿里巴巴集团": {
        "阿里云": "移动云",
        "钉钉": "云视讯",
        "无影云桌面": "云电脑",
        "阿里云盘": "移动云盘",
        "通义千问大模型": "九天大模型",
        "斑马智行": "移动车家"
    },
    "腾讯控股集团": {
        "腾讯会议": "云视讯",
        "腾讯云": "移动云",
        "腾讯云视立方": "移动高清",
        "QQ邮箱": "139邮箱",
        "混元大模型": "九天大模型"
    },
    "华为技术有限公司": {
        "华为云": "移动云",
        "华为云会议": "云视讯",
        "好望云服务": "移动看家",
        "华为云桌面": "云电脑",
        "华为云专线": "e企组网",
        "盘古大模型": "九天大模型",
        "华为车联网": "行车卫士"
    },
    "中国电信集团": {
        "天翼云": "移动云",
        "天翼看家": "移动看家",
        "天翼对讲": "和对讲",
        "天翼云盘": "移动云盘",
        "云堤": "专线卫士",
        "天翼云会议": "云视讯"
    },
    "中国联合网络通信集团": {
        "联通云": "移动云",
        "联通云联网": "e企组网",
        "联通物联": "千里眼",
        "联通云盾": "专线卫士"
    },
    "网易网络有限公司": {
        "网易云盘": "移动云盘",
        "网易邮箱": "139邮箱"
    },
    "百度公司": {
        "百度云": "移动云",
        "百度网盘": "移动云盘",
        "文心一言": "九天大模型"
    },
    "京东集团": {
        "京东云": "移动云",
        "京东云盘": "移动云盘"
    }
}

# 定义行业通用推荐映射
industry_recommendations = {
    "通信": ["移动云", "云视讯", "e企组网", "专线卫士"],
    "科技": ["移动云", "云电脑", "九天大模型", "移动云盘", "139邮箱"],
    "金融": ["专线卫士", "移动云", "集团语音", "移动ICT"],
    "能源": ["千里眼", "移动云", "行车卫士", "和对讲"],
    "基建": ["千里眼", "和对讲", "移动看家", "e企组网"],
    "教育": ["云视讯", "移动云盘", "移动高清", "云电脑"],
    "物流": ["行车卫士", "移动车家", "和对讲", "千里眼"],
    "医疗": ["移动康养", "移动云", "云视讯", "移动看家"],
    "制造": ["千里眼", "e企组网", "移动云", "行车卫士"],
    "餐饮": ["移动看家", "移动高清", "移动云盘", "云视讯"],
    "政务": ["专线卫士", "移动云", "云视讯", "移动ICT"],
    "零售": ["移动看家", "移动高清", "移动云盘", "e企组网"]
}

# 生成新闻标题和关键词
news_keywords = [
    ("数字化转型加速", "数字化转型"),
    ("5G应用蓬勃发展", "5G应用"),
    ("云计算成为企业标配", "云计算"),
    ("人工智能赋能千行百业", "人工智能"),
    ("智慧医疗解决看病难", "智慧医疗"),
    ("在线教育促进公平", "在线教育"),
    ("车联网技术突破", "车联网"),
    ("网络安全形势严峻", "网络安全"),
    ("远程办公成为新常态", "远程办公"),
    ("智慧物流降本增效", "智慧物流"),
    ("国产化替代加速", "国产化替代"),
    ("大模型技术竞争激烈", "大模型"),
    ("智能家居市场扩大", "智能家居"),
    ("企业上云需求增长", "企业上云"),
    ("数据安全法规完善", "数据安全"),
    ("餐饮行业数字化转型", "餐饮数字化"),
    ("政务云建设加速", "政务云"),
    ("零售行业智能化升级", "智能零售"),
    ("智慧城市建设推进", "智慧城市"),
    ("工业互联网发展迅速", "工业互联网")
]

# 生成三元组
triplets = []

# 1. 生成新闻提及关键词
for news, keyword in news_keywords:
    triplets.append([f"《{news}》新闻", "提及", keyword])

# 2. 生成新闻提及企业
# 确保每个行业至少有一个企业被新闻提及
used_companies = set()
for i, (news, keyword) in enumerate(news_keywords):
    # 根据关键词选择相关行业
    industry_options = []
    if "数字" in keyword or "云" in keyword or "智能" in keyword:
        industry_options = ["科技", "制造", "通信"]
    elif "医疗" in keyword:
        industry_options = ["医疗"]
    elif "教育" in keyword:
        industry_options = ["教育"]
    elif "物流" in keyword:
        industry_options = ["物流"]
    elif "安全" in keyword:
        industry_options = ["金融", "能源", "科技", "政务"]
    elif "车" in keyword:
        industry_options = ["制造", "物流"]
    elif "餐饮" in keyword:
        industry_options = ["餐饮"]
    elif "政务" in keyword:
        industry_options = ["政务"]
    elif "零售" in keyword:
        industry_options = ["零售"]
    else:
        industry_options = list(companies_by_industry.keys())

    # 随机选择一个行业
    industry = random.choice(industry_options)

    # 选择该行业的一个企业，优先选择尚未被提及的企业
    companies = companies_by_industry[industry]
    available_companies = [c for c in companies if c not in used_companies]

    if available_companies:
        company = random.choice(available_companies)
    else:
        company = random.choice(companies)

    used_companies.add(company)
    triplets.append([f"《{news}》新闻", "提及", company])

# 3. 生成企业属于行业的关系
for industry, companies in companies_by_industry.items():
    for company in companies:
        triplets.append([company, "属于", industry])

# 4. 生成企业使用竞品的关系
for company, competitors in competitor_mapping.items():
    for competitor_product, _ in competitors.items():
        triplets.append([company, "使用", competitor_product])

# 对于没有竞品信息的企业，添加一些通用竞品使用关系
other_companies = []
for companies in companies_by_industry.values():
    other_companies.extend(companies)

other_companies = list(set(other_companies) - set(competitor_mapping.keys()))

generic_competitors = {
    "金融": ["国外安全软件", "传统金融系统", "国外支付系统"],
    "能源": ["传统监控系统", "国外数据分析平台", "传统能源管理系统"],
    "基建": ["普通监控设备", "传统通信系统", "传统工程管理软件"],
    "教育": ["国外教学平台", "传统视频系统", "国外在线考试系统"],
    "医疗": ["传统医疗信息系统", "国外远程诊疗平台", "传统医院管理系统"],
    "制造": ["传统生产管理系统", "国外工业软件", "传统质量控制系统"],
    "物流": ["普通GPS系统", "传统调度平台", "国外物流管理系统"],
    "科技": ["国外云服务", "传统开发工具", "国外大数据平台"],
    "通信": ["传统网络设备", "国外通信解决方案", "传统基站设备"],
    "餐饮": ["传统点餐系统", "国外餐饮管理软件", "传统库存管理系统"],
    "政务": ["传统政务系统", "国外政务服务软件", "传统公文处理系统"],
    "零售": ["传统POS系统", "国外零售管理软件", "传统库存管理系统"]
}

for company in other_companies:
    # 找到公司所属行业
    industry = None
    for ind, companies in companies_by_industry.items():
        if company in companies:
            industry = ind
            break

    if industry and industry in generic_competitors:
        competitors = generic_competitors[industry]
        for competitor in competitors[:random.randint(1, 3)]:  # 每个公司添加1-3个竞品
            triplets.append([company, "使用", competitor])

# 5. 生成企业推荐使用中国移动产品的关系
# 对于有竞品信息的企业，根据竞品映射推荐
for company, competitors in competitor_mapping.items():
    for competitor_product, cmcc_product in competitors.items():
        triplets.append([company, "（推荐）使用", cmcc_product])

# 对于其他企业，根据行业推荐
for company in other_companies:
    # 找到公司所属行业
    industry = None
    for ind, companies in companies_by_industry.items():
        if company in companies:
            industry = ind
            break

    if industry and industry in industry_recommendations:
        recommendations = industry_recommendations[industry]
        for product in recommendations[:random.randint(2, 4)]:  # 每个公司推荐2-4个产品
            triplets.append([company, "（推荐）使用", product])

# 6. 如果需要更多记录，添加额外的新闻提及关系
if len(triplets) < TARGET_RECORDS:
    additional_needed = TARGET_RECORDS - len(triplets)

    # 创建更多新闻
    additional_news = [
        ("工业互联网平台建设", "工业互联网"),
        ("智慧农业助力乡村振兴", "智慧农业"),
        ("数字人民币试点扩大", "数字人民币"),
        ("元宇宙技术发展", "元宇宙"),
        ("区块链应用拓展", "区块链"),
        ("边缘计算需求增长", "边缘计算"),
        ("物联网设备普及", "物联网"),
        ("智慧交通缓解拥堵", "智慧交通"),
        ("远程医疗技术成熟", "远程医疗"),
        ("数字孪生技术应用", "数字孪生"),
        ("智能客服替代人工", "智能客服"),
        ("AR/VR技术商用", "AR/VR"),
        ("大数据分析提升效率", "大数据"),
        ("量子计算研究突破", "量子计算"),
        ("自动驾驶技术测试", "自动驾驶")
    ]

    # 添加新闻提及关键词
    for i in range(min(additional_needed // 2, len(additional_news))):
        news, keyword = additional_news[i]
        triplets.append([f"《{news}》新闻", "提及", keyword])
        additional_needed -= 1

    # 添加新闻提及企业
    all_companies = []
    for companies in companies_by_industry.values():
        all_companies.extend(companies)
    all_companies = list(set(all_companies))

    for i in range(min(additional_needed, len(all_companies))):
        company = all_companies[i % len(all_companies)]
        triplets.append([f"《{additional_news[i % len(additional_news)][0]}》新闻", "提及", company])
        additional_needed -= 1

    # 如果还需要更多记录，添加额外的推荐关系
    if additional_needed > 0:
        for i in range(additional_needed):
            company = random.choice(all_companies)
            product = random.choice(cmcc_products)
            triplets.append([company, "（推荐）使用", product])

# 确保不超过目标记录数量
if len(triplets) > TARGET_RECORDS:
    triplets = triplets[:TARGET_RECORDS]

# 创建数据目录（如果不存在）
os.makedirs('./data', exist_ok=True)

# 写入CSV文件
with open('./data/gov_bus_knowledge_graph.csv', 'w', newline='', encoding='utf-8') as csvfile:
    writer = csv.writer(csvfile)
    writer.writerow(['head', 'relation', 'tail'])  # 写入表头
    writer.writerows(triplets)

print(f"知识图谱CSV文件已生成: ./data/gov_bus_knowledge_graph.csv")
print(f"总共生成 {len(triplets)} 条三元组记录")
print(f"涵盖 {len(companies_by_industry)} 个行业, {sum(len(v) for v in companies_by_industry.values())} 家企业")
print(f"涵盖 {len(cmcc_products)} 种中国移动产品")
